host_build_graph: parallelize the boot initial-classify across AICPU threads - #1553
Merged
ChaoZheng109 merged 1 commit intoJul 29, 2026
Merged
Conversation
…threads The device boot classify — seed the whole graph's ready queues and wake lists (route roots to the ready queues, register every other task on its first unmet producer's wake list) — used to run O(total_tasks) serially on the boot leader while the other AICPU threads idle-waited on runtime_init_ready_. On a 65792-task paged_attention graph that is ~2.26 ms of one thread working and three spinning. Split it: on_orchestration_done keeps only the leader-only setup (attach, task count, queue allocation) and the classify moves to classify_partition(), which each thread runs over a disjoint contiguous slice of the submitted-task range. The boot leader publishes classify_ready_ once its setup is visible; every thread then classifies its slice and increments the classify_arrived_ barrier; the leader publishes runtime_init_ready_ only once all slices are done, so no thread dispatches against a half-seeded graph. push_ready_routed and register_wake are the same lock-free primitives the scheduler threads already use concurrently, and at boot no producer has completed (wake-list heads are nullptr, never SENTINEL), so registration never re-classifies — the parallel scan needs no new locking. Measured on a2a3, paged_attention batch=256 (65792 tasks), 4 AICPU threads: the classify phase drops from 2.26 ms (serial) to 1.58 ms (-30%). It does not reach the 4x of perfect scaling because concurrent register_wake on a high-fanout producer's wake list and push_ready_routed on the shared ready queues re-introduce the CAS contention the run-time 3S+1P split removes, plus the barrier waits on the slowest slice — but it still returns ~0.68 ms of otherwise-idle serial boot time to the run. matmul, bgemm, paged_attention, vector_example and predicated_dispatch all pass (the barrier introduces no deadlock and the seeded graph is correct).
📝 WalkthroughWalkthroughAICPU startup now partitions initial task classification across threads. Each thread seeds ready queues or wake lists for its task range, and runtime initialization becomes ready only after all classification partitions complete. ChangesParallel initial classification
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant BootLeader
participant AICPUThreads
participant SchedulerContext
BootLeader->>AICPUThreads: Publish classify_ready_
AICPUThreads->>SchedulerContext: Classify each thread partition
SchedulerContext-->>AICPUThreads: Seed ready queues or wake lists
AICPUThreads->>AICPUThreads: Wait for all classify arrivals
AICPUThreads-->>BootLeader: Publish runtime_init_ready_
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp`:
- Around line 320-329: Update the worker startup/classification barrier around
classify_arrived_ and runtime_init_ready_ to use a shared abort state when an
out-of-range worker exits before arrival. Ensure every boot wait checks this
state, and change the early invalid-affinity return to signal the abort and
follow normal shutdown so valid workers cannot spin indefinitely.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3afc291e-b710-4e58-98d3-fdbbf301df33
📒 Files selected for processing (3)
src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Parallelize the
host_build_graphdevice boot initial-classify across the AICPU threads.The boot classify seeds the whole graph's ready queues and wake lists — route roots (all fanin already met) to the ready queues, register every other task on its first unmet producer's wake list. It ran
O(total_tasks)serially on the boot leader while the other AICPU threads idle-waited onruntime_init_ready_.This PR splits it:
on_orchestration_donekeeps only the leader-only setup (attach, task-count latch, SPSC queue allocation).classify_partition(), which each thread runs over a disjoint contiguous slice of the submitted-task range.classify_ready_once its setup is visible; every thread classifies its slice and increments theclassify_arrived_barrier; the leader publishesruntime_init_ready_only once all slices finish — so no thread dispatches against a half-seeded graph.push_ready_routedandregister_wakeare the same lock-free primitives the scheduler threads already use concurrently at run time, and at boot no producer has completed (wake-list heads arenullptr, neverSENTINEL), so registration never re-classifies — the parallel scan needs no new locking, only the boot barrier (modeled on the existing handshake barrier).Scope: a2a3
host_build_graphonly. Follow-up to #1544.Why
On a 65792-task
paged_attentiongraph the serial classify is ~2.26 ms of one thread working while three spin — pure serial boot latency on the critical path before any task dispatches. The other threads are right there, idle-waiting, and the classify is embarrassingly parallel (each task is independent).Performance (a2a3,
paged_attentionCase1, 65792 tasks, 4 AICPU threads)Classify-phase wall (boot leader, from setup-done to
runtime_init_ready_):It does not reach the 4× of perfect scaling: concurrent
register_wakeon a high-fanout producer's wake list andpush_ready_routedon the shared ready queues re-introduce the CAS contention the run-time 3S+1P split removes, and the barrier waits on the slowest slice. The gain is workload-dependent (largest on scheduler-bound graphs with many small tasks) — ~0.68 ms is ~2.7% of this run's ~25 ms device execution, and negligible for compute-bound long runs. Still a strict improvement: otherwise-idle serial boot time handed back, off the critical-path start.Correctness
matmul,bgemm,paged_attention(Case1),vector_exampleandpredicated_dispatchall pass on a2a3 — the barrier introduces no deadlock and the parallel-seeded ready-set / wake-list graph is correct.🤖 Generated with Claude Code